#!/usr/bin/env python

from numpy import linspace,zeros,dot
from scipy.linalg import solve
from collocation import lhs, rhs
from bernstein import interpolant
from jacobian import evaluate_fun, evaluate_jac
from exact import exact_solution
from matplotlib import pyplot as plt

n = 60  # Bernstein polynomial order

a = 0.  # Solution interval [a,b]
b = 0.9

niter = 7  # Number of Newton iterations
tol = 1e-5 # Criteria to exit Newton iterations. If norm of update vector < tol, exit.

beta = zeros(n+1)  # Bernstein polynomial coefficients - the principal unknown
# Boundary conditions - direct imposition
beta[0] = 1.       # Here you put f(a)     
beta[n] = 32.725   # Here you put f(b)

nd = linspace(a,b,n+1) # nodes

# Initialize solution for Newton method - a straight line from y(a) to y(b):
for i in range(1,n) :
	x = nd[i]
	dist = (x-a)/(b-a)
	beta[i] = beta[0] + dist*(beta[n]-beta[0])

# Make initial solution:
uin = zeros(n+1)
for i in range(n+1):
	x = nd[i]
	uin[i] = interpolant(beta,n,a,b,x)


### NEWTON's ITERATIONS ###
#     We solve in following steps using Newton's algorithm:
#     1. Solve linear system: DF(xi) * Dx = -f(xi); DF is Jacobian matrix of nl. func. F
#     2. x(i+1) = x(i) + Dx

delta = zeros(n-1)

for it in range(1,niter+1):

	f = zeros(n-1)
	for i in range(n-1):
		x = nd[i+1]
		f[i] = - evaluate_fun(x,a,b,n,beta) # RHS vector entry

	K = zeros( (n-1,n-1) )
	for i in range(n-1):
		x = nd[i+1]
		for j in range(n-1):
			K[i,j] = evaluate_jac(x,j+1,a,b,n,beta) # LHS matrix entry

	delta = solve(K,f)  # Solve linear system

	beta[1:n] = beta[1:n] + delta[0:n-1]  # Update solution

        nrm = dot(delta,delta)
        if (nrm < tol): break 

### END: NEWTON's ITERATIONS ###

u = zeros(n+1)
for i in range(n+1):
	x = nd[i]
	u[i] = interpolant(beta,n,a,b,x)  # Solution - a Bernstein interpolant constructed using beta's
                                      # at whatever points - we use collocation points here.

# Print absolute error
print 'Node, Initial solution, and collocation solution after',niter,' Newton iterations.'
print u

plt.plot(nd,uin,'r', nd, u, 'g')
plt.legend(('Initial sol.', 'Final sol.'), loc='upper left')
plt.show()



